import torch
from torch import Tensor
import numpy as np

from torchsponge.system import Molecule
from torchsponge.optimizer import md

from torchsponge.potential.pyscfcell import PySCFForceCell

from torchsponge.core import WithForceCell, RunOneStepCell
from torchsponge.optimizer import UpdaterMD
from torchsponge.partition import NeighbourList
from torchsponge.control import Langevin, LeapFrog
from torchsponge.core.sponge import Sponge
from torchsponge.callback import WriteH5MD
from torchsponge.function import VelocityGenerator

coord = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.074]], dtype=float)
atom_types = np.array([1,1],dtype=int)
mol = Molecule(
    atomic_number=atom_types,
    coordinate=coord
)

pyscf_frc  = PySCFForceCell(system=mol, pyscf_option='ks', mol_parameters={'basis': '6-31g','verbose':0}, ks_parameters={'xc':'b3lyp'}, use_pbc=False)
wf = WithForceCell(mol, pyscf_frc)

velocity = torch.zeros_like(mol.coordinate)
opt = UpdaterMD(mol, 1e-3, velocity, 0, None, 'leap_frog')
onestep = RunOneStepCell(force=wf, optimizer=opt)
md = Sponge(network=onestep, write_h5md=WriteH5MD(mol, 'test.h5md', 100, write_velocity=True, write_force=True))

md.run(10,1,1)  # total steps, print_interval, write_info_interval